"""
3.	请用深度学习平台，实现yolov3的主干网络darknet53（参照下图），完成CIFAR10数据的10分类（25分）
"""
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers, activations, optimizers, losses, metrics
from sklearn.model_selection import train_test_split
import numpy as np

ALPHA = 1e-3
BATCH_SIZE = 32
N_EPOCH = 2


# ②	定义类封装卷积单元ConvCell，包含卷积、BN、LeekyRelu
class ConvCell(keras.Model):

    def __init__(self, filters, ksize=(3, 3), strides=(1, 1), padding='same'):
        super().__init__()
        self.conv = layers.Conv2D(filters, ksize, strides, padding, use_bias=False)
        self.bn = layers.BatchNormalization()
        self.lrelu = layers.LeakyReLU()

    def call(self, inputs, training=None, mask=None):
        x = self.conv(inputs, training=training)
        x = self.bn(x, training=training)
        x = self.lrelu(x, training=training)
        return x


# ③	定义类封装卷积模块ConvBlock，包含两个卷积单元（第一个卷积单元的核是1x1，第二个卷积单元的核是3x3）和一个残差模块
class ConvBlock(keras.Model):

    def __init__(self, filters):
        super().__init__()
        self.conv1 = ConvCell(filters // 2, (1, 1))
        self.conv2 = ConvCell(filters)

    def call(self, inputs, training=None, mask=None):
        x = self.conv1(inputs, training=training)
        x = self.conv2(x, training=training)
        x += inputs
        return x


# ④	定义类封装yolov3主干模型Darknet53
class Darknet53(keras.Model):

    def __init__(self, n_cls):
        super().__init__()
        filters = 32
        self.conv1 = ConvCell(filters)
        cfg = [1, 2, 8, 8, 4]
        seqs = []
        for n_block in cfg:
            filters *= 2
            seqs.append(ConvCell(filters, (3, 3), (2, 2)))
            for i in range(n_block):
                seqs.append(ConvBlock(filters))
        self.main = keras.Sequential(seqs)
        self.avg = layers.GlobalAvgPool2D()
        self.fc = layers.Dense(n_cls, activation=activations.softmax)

    def call(self, inputs, training=None, mask=None):
        x = self.conv1(inputs, training=training)
        x = self.main(x, training=training)
        x = self.avg(x, training=training)
        x = self.fc(x, training=training)
        return x


# ⑤	打印输出Darknet53模型结构维度
model = Darknet53(10)
model.build(input_shape=(None, 32, 32, 3))
model.summary()

# ①	导入CIFAR10的测试集作为模型训练的数据集，然后再按适当比例划分训练集、验证集、测试集
(_, _), (x, y) = keras.datasets.cifar10.load_data()
x = np.float32(x) / 255.
x_train, x_val_test, y_train, y_val_test = train_test_split(x, y, train_size=0.8, random_state=1, shuffle=True)
x_val, x_test, y_val, y_test = train_test_split(x_val_test, y_val_test, train_size=0.5, random_state=1, shuffle=True)
print('x_train', x_train.shape)
print('x_val', x_val.shape)
print('x_test', x_test.shape)
print('y_train', y_train.shape)
print('y_val', y_val.shape)
print('y_test', y_test.shape)

# ⑥	进行模型编译和训练，打印输出训练集、验证集、测试集准确率
model.compile(
    optimizer=optimizers.Adam(learning_rate=ALPHA),
    loss=losses.sparse_categorical_crossentropy,
    metrics=[metrics.sparse_categorical_accuracy]
)
model.fit(
    x_train,
    y_train,
    batch_size=BATCH_SIZE,
    epochs=N_EPOCH,
    validation_data=(x_val, y_val)
)
print('Testing ...')
model.evaluate(x_test, y_test, BATCH_SIZE)
print('Tested')
